package main

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"runtime"
	"shell/common"
	"shell/common/packages"
	"shell/controlled/config"
	"shell/global"
	"shell/utils"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/dnephin/pflag"
	"github.com/google/uuid"
	"github.com/quic-go/quic-go"
)

var (
	argGenUUID                            = pflag.BoolP("uuid", "u", false, "generate a uuid and exit")
	argDevMode                            = pflag.BoolP("dev-mode", "d", false, "develop mode, print groutine number every 10 seconeds")
	argCfgPath                            = pflag.StringP("config", "c", "./config.yaml", "set config file path")
	bytePool                              = utils.NewBytesPool()
	globalCfg      *config.ControlledCfg  = nil
	logger                                = global.DefaultLogger()
	timeoutWatcher *common.TimeoutWatcher = nil
)

func main() {
	// 解析参数
	pflag.Parse()

	// 生成uuid
	if *argGenUUID {
		fmt.Printf("uuid: %s \n", uuid.New())
		os.Exit(0)
	}

	if *argDevMode {
		logger.SetLevel(global.Debug)
	}

	// 解析配置
	cfg, err := config.ParseConfig(*argCfgPath)
	if err != nil {
		logger.Fatalf(1, "error read config file: %v", err)
	}
	globalCfg = cfg

	// globalCtx为全局上下文，调用globalCancel会时mianLogic方法退出
	globalCtx, globalCancel := context.WithCancel(context.Background())
	// 创建超时监控器
	timeoutWatcher = common.NewTimeoutWatcher(globalCtx, 30*time.Second)

	// 处理系统信号量
	sigs := make(chan os.Signal, 1)
	connChan := make(chan *common.QuicConn, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	getSysSem := atomic.Bool{}
	getSysSem.Store(false)
	wg := sync.WaitGroup{}
	go func() {
		var currentConn *common.QuicConn
		for {
			select {
			case conn := <-connChan:
				currentConn = conn
			case s := <-sigs:
				wg.Add(1)
				getSysSem.Store(true)
				logger.Infof(0, "get operating system semaphores: %s, exit", s)
				globalCancel()
				stream, err := currentConn.OpenStream()
				if err != nil {
					common.HandleErr(err)
					stream.CloseConn(common.CloseByProcess)
					wg.Done()
					return
				}
				_ = stream.CloseConnSync(Role, true)
				wg.Done()
				return
			case <-globalCtx.Done():
				return
			}
		}
	}()

	for {
		start := time.Now()
		// 创建quic连接
		conn, err := quic.DialAddr(context.Background(), globalCfg.ProxyAddr, globalCfg.TLS, &quic.Config{
			MaxIdleTimeout:     time.Second * 15,
			MaxIncomingStreams: 1 << 40,
		})
		if err != nil {
			logger.Fatalf(0, "error dial to proxy: %v", err)
		}
		quicConn := common.NewQuicConn(conn, bytePool)
		connChan <- quicConn
		// 启动主逻辑流程
		mainLogic(quicConn, globalCtx)
		keepDuration := time.Since(start)
		if getSysSem.Load() {
			// 等待信号量goroutine结束
			wg.Wait()
			break
		}
		if keepDuration < time.Second*15 {
			logger.Errorf(0, "error main logic exit too quick, duration: %v", keepDuration)
			break
		}
		time.Sleep(time.Second * 15)
	}
	globalCancel()
}

// 主逻辑方法
func mainLogic(quicConn *common.QuicConn, ctx context.Context) {
	go keepalive(quicConn, ctx)
	// 处理quic连接，阻塞
	handleConn(quicConn, ctx)
}

// 执行保活命令
func keepalive(quicConn *common.QuicConn, ctx context.Context) {
	stream, err := quicConn.OpenStream()
	if err != nil {
		common.HandleErr(err)
		return
	}
	err = keepaliveOnce(stream)
	if err != nil {
		common.HandleErr(err)
		quicConn.Close(common.CloseByProcess)
		return
	}
	req := new(packages.KeepaliveReq)
	req.SetRole(Role)
	ticker := time.NewTicker(time.Second * 6)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			if *argDevMode {
				logger.Debugf(0, "goroutine num: %d", runtime.NumGoroutine())
			}
			stream, err := quicConn.OpenStream()
			if err != nil {
				common.HandleErr(err)
				return
			}
			err = keepaliveOnce(stream)
			if err != nil {
				common.HandleErr(err)
				quicConn.Close(common.CloseByProcess)
				return
			}
		}
	}
}

// 执行一次保活请求并接收响应，如果发生错误，就关闭quic连接
// 如果没有错误，会关闭流
func keepaliveOnce(stream *common.QuicStream) error {
	req := new(packages.KeepaliveReq)
	req.SetRole(Role)
	req.CmdUUID = uuid.New()
	info := globalCfg.ToKeepaliveInfo()
	i, _, _ := utils.ParseIpPort(globalCfg.ProxyAddr)
	ip := utils.GetIPFrom(i)
	if ip == nil {
		stream.CloseConn(common.CloseByProcess)

		return errors.New("invalid ip")
	}
	info.PrivateIP = ip
	info.UUID = globalCfg.UUID
	req.Stats = []*packages.KeepaliveInfo{info}
	err := stream.WritePackage(req)
	if err != nil {
		stream.CloseConn(common.CloseByProcess)
		return err
	}
	resp := new(packages.KeepaliveResp)
	err = stream.ReadPackage(resp)
	if err != nil {
		stream.CloseConn(common.CloseByProcess)
		return err
	}
	if resp.CmdUUID != req.CmdUUID {
		stream.CloseConn(common.CloseByProcess)
		return packages.ErrCmdUUIDChanged
	}
	_ = stream.CloseStreamSync(Role, true)
	return nil
}

// 处理请求
func handleConn(conn *common.QuicConn, ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
			stream, err := conn.AcceptStream()
			if err != nil {
				common.HandleErr(err)
				return
			}
			go handleStream(stream)
		}
	}
}

// 处理流，完整接收请求，根据请求类型处理请求
func handleStream(stream *common.QuicStream) {
	defer func() {
		_ = (*stream).Close()
	}()
	md, err := common.DetectPackage(stream)
	if err != nil {
		common.HandleErr(err)
		return
	}
	pkgType := md.PkgAct
	switch pkgType {
	case packages.ActSocks5RedirReq:
		// sock5请求
		pkg := new(packages.Socks5RedirReq)
		err = stream.ReadPackageAfterDetect(pkg, md)
		if err != nil {
			common.HandleErr(err)
			return
		}
		handleSocks5(stream, pkg)
	case packages.ActPortRedirReq:
		// 端口转发请求
		pkg := new(packages.PortRedirReq)
		err = stream.ReadPackageAfterDetect(pkg, md)
		if err != nil {
			common.HandleErr(err)
			return
		}
		handlePortRedir(stream, pkg)
	case packages.ActShellCmdReq:
		// 执行shell命令请求
		pkg := new(packages.ShellCmdReq)
		err = stream.ReadPackageAfterDetect(pkg, md)
		if err != nil {
			common.HandleErr(err)
			return
		}
		handleShellCmd(stream, pkg)
	case packages.ActKeepaliveReq:
		// 由代理端到被控端的保活
		req := new(packages.KeepaliveReq)
		err = stream.ReadPackageAfterDetect(req, md)
		if err != nil {
			common.HandleErr(err)
			return
		}
		// 校验保活请求是否来自代理端
		if req.GetRole() != packages.RoleProxy {
			logger.Errorf(0, "error receive keep req from: %s", req.GetRole().String())
		}
		resp := new(packages.KeepaliveResp)
		resp.SetRole(Role)
		resp.CmdUUID = req.CmdUUID
		err = stream.WritePackage(resp)
		if err != nil {
			common.HandleErr(err)
			return
		}
		stream.WaitForClose(Role, false)
	default:
		// 不支持的请求类型
		logger.Errorf(0, "error package action: %d", pkgType)
		_ = stream.Close()
		return
	}
}
